#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <pthread.h>
#include "defs.h"
#include "common_threads.h"

typedef struct {
    int start_row;
    int end_row;
    int start_col;
    int end_col;
    double *a;
    double *b;
    double *c;
    int m;
    int n;
    int k;
    int lda;
    int ldb;
    int ldc;
} myarg_t;


void *matrix_multiply_block(void *arg) {
    myarg_t *args = (myarg_t *) arg;
    myarg_t *rvals = malloc(sizeof(myarg_t));
    int start_row = args->start_row;
    int end_row = args->end_row;
    int start_col = args->start_col;
    int end_col = args->end_col;
    double *a = args->a;
    double *b = args->b;
    double *c = args->c;
    int m = args->m;
    int n = args->n;
    int k = args->k;
    int lda = args->lda;
    int ldb = args->ldb;
    int ldc = args->ldc;
    rvals->c = c;

    for (int i = start_row; i < end_row; i++) 
        for (int j = start_col; j < end_col; j++) 
            for (int p = 0; p < k; p++) 
                C(i, j) += A(i, p) * B(p, j);
            
    return (void *) rvals;
}

void MY_MMult(int m, int n, int k, double *a, int lda,
              double *b, int ldb,
              double *c, int ldc) {

    int num_threads = 8;
    pthread_t threads[num_threads];
    myarg_t args[num_threads];
    myarg_t *rvals[num_threads];

    int block_size_row = m / 4;
    int block_size_col = n / 2;

    for (int i = 0; i < num_threads; i++) {
        args[i].start_row = (i % 4) * block_size_row;
        args[i].end_row = (args[i].start_row) + block_size_row;
        args[i].start_col = (i / 4) * block_size_col;
        args[i].end_col = (args[i].start_col) + block_size_col;
        args[i].a = a;
        args[i].b = b;
        args[i].c = c;
        args[i].m = m;
        args[i].n = n;
        args[i].k = k;
        args[i].lda = lda;
        args[i].ldb = ldb;
        args[i].ldc = ldc;

        Pthread_create(&threads[i], NULL, matrix_multiply_block, &args[i]);
    }

    for (int i = 0; i < num_threads; i++) {
        Pthread_join(threads[i], (void **) &rvals[i]);
    }
    for (int i = 0; i < num_threads; i++)
        free(rvals[i]);
}
